#!/usr/bin/env python
"""
清理旧的checkpoint，只保留最佳模型和间隔checkpoint
使用方法: python cleanup_checkpoints.py --checkpoint_dir ./model/uncond_checkpoints_t3 --keep_interval 10000
"""
import os
import glob
import argparse
import re
from collections import defaultdict


def parse_checkpoint_name(filename):
    """解析checkpoint文件名，提取step number"""
    # 匹配 model000000.pt, ema_0.9999_000000.pt, opt000000.pt
    patterns = [
        (r'model(\d+)\.pt', 'model'),
        (r'ema_[\d\.]+_(\d+)\.pt', 'ema'),
        (r'opt(\d+)\.pt', 'opt'),
    ]
    
    for pattern, ckpt_type in patterns:
        match = re.search(pattern, filename)
        if match:
            step = int(match.group(1))
            return step, ckpt_type
    return None, None


def get_checkpoint_groups(checkpoint_dir):
    """将checkpoint按step分组"""
    checkpoint_files = glob.glob(os.path.join(checkpoint_dir, "*.pt"))
    
    groups = defaultdict(list)
    best_files = []
    
    for filepath in checkpoint_files:
        filename = os.path.basename(filepath)
        
        # 跳过best checkpoint
        if filename.startswith("best_"):
            best_files.append(filepath)
            continue
        
        step, ckpt_type = parse_checkpoint_name(filename)
        if step is not None:
            groups[step].append(filepath)
    
    return groups, best_files


def cleanup_checkpoints(checkpoint_dir, keep_interval=10000, dry_run=False, keep_latest_n=2):
    """
    清理checkpoint，保留:
    1. 所有best_开头的文件
    2. 间隔点的checkpoint (step % keep_interval == 0)
    3. 最新的N个checkpoint
    
    Args:
        checkpoint_dir: checkpoint目录
        keep_interval: 保留间隔（如10000）
        dry_run: 是否只模拟不实际删除
        keep_latest_n: 保留最新的N个checkpoint
    """
    print(f"扫描目录: {checkpoint_dir}")
    print(f"保留间隔: 每 {keep_interval} 步")
    print(f"保留最新: {keep_latest_n} 个checkpoint")
    print(f"模式: {'模拟运行' if dry_run else '实际删除'}")
    print("-" * 60)
    
    groups, best_files = get_checkpoint_groups(checkpoint_dir)
    
    if not groups:
        print("未找到checkpoint文件")
        return
    
    print(f"\n找到 {len(best_files)} 个best checkpoint (将全部保留)")
    print(f"找到 {len(groups)} 个不同step的checkpoint组")
    
    # 按step排序
    sorted_steps = sorted(groups.keys())
    latest_steps = sorted_steps[-keep_latest_n:]
    
    print(f"\n最新的 {keep_latest_n} 个checkpoint: {latest_steps}")
    
    # 统计
    total_files = sum(len(files) for files in groups.values())
    to_remove = []
    to_keep = []
    
    for step in sorted_steps:
        files = groups[step]
        
        # 判断是否应该保留
        should_keep = False
        reason = ""
        
        if step % keep_interval == 0:
            should_keep = True
            reason = f"间隔点 (step % {keep_interval} == 0)"
        elif step in latest_steps:
            should_keep = True
            reason = f"最新的 {keep_latest_n} 个之一"
        
        if should_keep:
            to_keep.extend(files)
            print(f"\n✓ 保留 step {step:6d} ({len(files)} 文件) - {reason}")
        else:
            to_remove.extend(files)
            print(f"\n✗ 删除 step {step:6d} ({len(files)} 文件)")
            if not dry_run:
                for filepath in files:
                    try:
                        os.remove(filepath)
                        print(f"    已删除: {os.path.basename(filepath)}")
                    except Exception as e:
                        print(f"    删除失败: {os.path.basename(filepath)} - {e}")
    
    # 总结
    print("\n" + "=" * 60)
    print(f"总计:")
    print(f"  Best checkpoint: {len(best_files)} 个文件 (已保留)")
    print(f"  保留的checkpoint: {len(to_keep)} 个文件")
    print(f"  {'将要' if dry_run else '已'}删除的checkpoint: {len(to_remove)} 个文件")
    
    if to_remove:
        # 计算节省的空间
        if dry_run:
            total_size = sum(os.path.getsize(f) for f in to_remove)
            print(f"  预计节省空间: {total_size / (1024**3):.2f} GB")
        else:
            print(f"  实际节省空间: {sum(os.path.getsize(f) for f in to_keep if os.path.exists(f)) / (1024**3):.2f} GB (保留文件总大小)")
    
    print("=" * 60)


def main():
    parser = argparse.ArgumentParser(description='清理旧的checkpoint文件')
    parser.add_argument('--checkpoint_dir', type=str, required=True,
                        help='Checkpoint目录路径')
    parser.add_argument('--keep_interval', type=int, default=10000,
                        help='保留间隔步数 (默认: 10000)')
    parser.add_argument('--keep_latest', type=int, default=2,
                        help='保留最新的N个checkpoint (默认: 2)')
    parser.add_argument('--dry_run', action='store_true',
                        help='只模拟不实际删除')
    parser.add_argument('--confirm', action='store_true',
                        help='跳过确认直接执行')
    
    args = parser.parse_args()
    
    if not os.path.exists(args.checkpoint_dir):
        print(f"错误: 目录不存在: {args.checkpoint_dir}")
        return
    
    # 显示信息
    print("\n" + "=" * 60)
    print("Checkpoint 清理工具")
    print("=" * 60)
    
    if not args.dry_run and not args.confirm:
        print("\n⚠️  警告: 此操作将永久删除文件!")
        print("建议先使用 --dry_run 模式查看将要删除的文件")
        response = input("\n确认继续? (yes/no): ")
        if response.lower() not in ['yes', 'y']:
            print("已取消")
            return
    
    cleanup_checkpoints(
        args.checkpoint_dir,
        keep_interval=args.keep_interval,
        dry_run=args.dry_run,
        keep_latest_n=args.keep_latest
    )


if __name__ == '__main__':
    main()

